import io
import base64
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_samples, silhouette_score
from sklearn.preprocessing import StandardScaler

# 聚类类别名称映射
CLUSTER_NAMES = {
    0: "高评分新片",      # 近期发行且评分较高的电影
    1: "经典老片",       # 发行时间较早但评分较高的电影
    2: "小众佳作",       # 评分不错但关注度不高的电影
    3: "热门商业片",     # 近期发行且关注度高的电影
    4: "待发掘佳作"      # 评分不错但发行时间较新的电影
}

def kmeans_visualization(dataset_id, n_clusters=3):
    """K-means聚类可视化"""
    try:
        # 读取数据
        df = pd.read_csv(f'uploads/{dataset_id}.csv')
        
        # 数据预处理
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        X = df[numeric_cols].fillna(0)
        
        # 标准化
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)
        
        # 执行K-means聚类
        kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        df['cluster'] = kmeans.fit_predict(X_scaled)
        
        # 计算轮廓系数
        silhouette_avg = silhouette_score(X_scaled, df['cluster'])
        
        # 1. 散点图可视化
        plt.figure(figsize=(12, 8))
        scatter = plt.scatter(X_scaled[:, 0], X_scaled[:, 1], c=df['cluster'], cmap='viridis')
        plt.title('K-means聚类结果')
        plt.xlabel('特征1')
        plt.ylabel('特征2')
        
        # 添加聚类中心
        centers = kmeans.cluster_centers_
        plt.scatter(centers[:, 0], centers[:, 1], c='red', marker='x', s=200, linewidths=3, label='聚类中心')
        
        # 添加图例
        legend_labels = [f'{CLUSTER_NAMES[i]}' for i in range(n_clusters)]
        plt.legend(handles=scatter.legend_elements()[0], labels=legend_labels, title="聚类类别")
        
        # 保存散点图
        scatter_buffer = io.BytesIO()
        plt.savefig(scatter_buffer, format='png', bbox_inches='tight')
        scatter_buffer.seek(0)
        scatter_base64 = base64.b64encode(scatter_buffer.getvalue()).decode()
        plt.close()
        
        # 2. 轮廓图
        plt.figure(figsize=(12, 8))
        silhouette_vals = silhouette_samples(X_scaled, df['cluster'])
        y_lower = 10
        
        for i in range(n_clusters):
            cluster_silhouette_vals = silhouette_vals[df['cluster'] == i]
            cluster_silhouette_vals.sort()
            
            size_cluster_i = len(cluster_silhouette_vals)
            y_upper = y_lower + size_cluster_i
            
            color = plt.cm.viridis(float(i) / n_clusters)
            plt.fill_betweenx(np.arange(y_lower, y_upper),
                             0, cluster_silhouette_vals,
                             facecolor=color, edgecolor=color, alpha=0.7)
            
            plt.text(-0.05, y_lower + 0.5 * size_cluster_i, CLUSTER_NAMES[i])
            y_lower = y_upper + 10
            
        plt.axvline(x=silhouette_avg, color="red", linestyle="--")
        plt.title('聚类轮廓分析图')
        plt.xlabel('轮廓系数')
        plt.ylabel('聚类')
        
        # 保存轮廓图
        silhouette_buffer = io.BytesIO()
        plt.savefig(silhouette_buffer, format='png', bbox_inches='tight')
        silhouette_buffer.seek(0)
        silhouette_base64 = base64.b64encode(silhouette_buffer.getvalue()).decode()
        plt.close()
        
        # 3. 热力图
        plt.figure(figsize=(12, 8))
        cluster_centers = pd.DataFrame(kmeans.cluster_centers_, columns=numeric_cols)
        cluster_centers.index = [CLUSTER_NAMES[i] for i in range(n_clusters)]
        sns.heatmap(cluster_centers, annot=True, cmap='YlOrRd', fmt='.2f')
        plt.title('聚类中心特征热力图')
        
        # 保存热力图
        heatmap_buffer = io.BytesIO()
        plt.savefig(heatmap_buffer, format='png', bbox_inches='tight')
        heatmap_buffer.seek(0)
        heatmap_base64 = base64.b64encode(heatmap_buffer.getvalue()).decode()
        plt.close()
        
        # 4. 雷达图
        plt.figure(figsize=(12, 8))
        # 选择前5个重要特征
        important_features = numeric_cols[:5]
        angles = np.linspace(0, 2*np.pi, len(important_features), endpoint=False)
        angles = np.concatenate((angles, [angles[0]]))  # 闭合图形
        
        ax = plt.subplot(111, polar=True)
        for i in range(n_clusters):
            values = cluster_centers.loc[CLUSTER_NAMES[i], important_features].values
            values = np.concatenate((values, [values[0]]))  # 闭合图形
            ax.plot(angles, values, 'o-', linewidth=2, label=CLUSTER_NAMES[i])
            ax.fill(angles, values, alpha=0.25)
            
        plt.xticks(angles[:-1], important_features)
        plt.title('聚类特征雷达图')
        plt.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1))
        
        # 保存雷达图
        radar_buffer = io.BytesIO()
        plt.savefig(radar_buffer, format='png', bbox_inches='tight')
        radar_buffer.seek(0)
        radar_base64 = base64.b64encode(radar_buffer.getvalue()).decode()
        plt.close()
        
        return {
            'scatter': scatter_base64,
            'silhouette': silhouette_base64,
            'heatmap': heatmap_base64,
            'radar': radar_base64
        }
        
    except Exception as e:
        raise Exception(f"可视化生成失败: {str(e)}") 